package com.feidee.fd.sml.algorithm.component.ml.classification

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.FFM

/**
  * @Author songhaicheng
  * @Date 2019/2/27 21:20
  * @Description
  * @Reviewer
  */
case class FFMParam (
                      override val input_pt: String,
                      override val output_pt: String,
                      override val hive_table: String,
                      override val flow_time: String,
                      override val featuresCol: String,
                      override var labelCol: String,
                      override var predictionCol: String,
                      override val modelPath: String,
                      override val metrics: Array[String],
                      // 原始预测值列名，默认 rawPrediction
                      rawPredictionCol: String,
                      // 预测概率列名，默认 probability
                      probabilityCol: String,
                      // 优化算法，支持 [sgd, adag]，默认 adag
                      optimizer: String,
                      // 迭代数，需大于 0，不填或填 0 时使用算法默认值
                      numIterations: Int,
                      // 隐向量长度，默认 2
                      k: Int,
                      // 分区数，大于等于 0，默认 0
                      partitions: Int,
                      // Spark treeAggregate 算子的参数，适度增大可以加快训练速度，大于等于 2，默认 2
                      aggregationDepth: Int,
                      // 随机种子，默认 123456
                      seed: Long,
                      // 采样率，每次选多少比例的样本构成新树，(0, 1]，默认 1.0
                      subsamplingRate: Double,
                      // 每次优化的迭代步长，> 0，默认 0.01
                      stepSize: Double,
                      // 正则化系数，>= 0，默认 0.0
                      regParam: Double,
                      // 收敛系数，>= 0，默认 1E-6
                      convergenceTol: Double,
                      // 是否对更新的权重进行正则化，默认 true
                      normalization: Boolean
                    ) extends MLParam {
  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    "rawPrediction", "probability",
    "adag", 200, 2, 0, 2, 123456, 1.0, 0.01, 0.0, 0.001, true)

  override def verify(): Unit = {
    super.verify()
    val optimizers = Array("sgd", "adag")
    require(optimizers.contains(optimizer.toLowerCase), s"param optimizer only accepts [${optimizers.mkString(", ")}]," +
      s" but has $optimizer")
    require(partitions >= 0, "param partitions can't be negative")
    require(aggregationDepth >= 2, "param aggregationDepth must be greater than 1")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("rawPredictionCol" -> rawPredictionCol)
    map += ("probabilityCol" -> probabilityCol)
    map += ("optimizer" -> optimizer)
    map += ("numIterations" -> numIterations)
    map += ("k" -> k)
    map += ("partitions" -> partitions)
    map += ("aggregationDepth" -> aggregationDepth)
    map += ("seed" -> seed)
    map += ("subsamplingRate" -> subsamplingRate)
    map += ("stepSize" -> stepSize)
    map += ("regParam" -> regParam)
    map += ("convergenceTol" -> convergenceTol)
    map += ("normalization" -> normalization)
    map
  }
}


class FFMComponent extends AbstractClassificationComponent[FFMParam] {

  override def setUp(param: FFMParam): PipelineStage = {
    val ffm = new FFM()
      .setLabelCol(param.labelCol)
      .setFeaturesCol(param.featuresCol)
      .setRawPredictionCol(param.rawPredictionCol)
      .setPredictionCol(param.predictionCol)
      .setProbabilityCol(param.probabilityCol)
      .setOptimizer(param.optimizer)
      .setMaxIter(param.numIterations)
      .setK(param.k)
      .setPartitions(param.partitions)
      .setAggregationDepth(param.aggregationDepth)
      //.setSeed(param.seed)
      .setMiniBatch(param.subsamplingRate)
      .setStepSize(param.stepSize)
      .setRegParam(param.regParam)
      .setConvergenceTol(param.convergenceTol)
      .setNormalization(param.normalization)

    ffm
  }

}

object FFMComponent {
  def apply(paramStr: String): Unit = {
    new FFMComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    FFMComponent(args(0))
  }
}
